#!/opt/local/bin/python

from rfvectors import cavity, scale_roverq
import matplotlib.pyplot as plt
from scipy import exp, nan, cos, log, sqrt, arctan, pi
from diffeqns import transientcavity, analysis, calcCurrents
rad2deg = lambda x: x*180/pi
import cProfile as profile

#sscavperf = cavity(f=704.42e6, delF=nan, I_b0=50e-3, I_b=50e-3, V_cav=18e6, beamph=-14.0)
#sscav     = cavity(f=704.42e6, delF=nan, I_b0=50e-3, I_b=50e-3, V_cav=18e6, beamph=-14.0, Ql=sscavperf.Ql)
sscavperf = cavity(f=704.21e6, delF=nan, I_b0=50e-3, I_b=50e-3, V_cav=18.7927e6*0.97885161, beamph=-14.0,
        RoverQ=477, Ql=8.123e5, beta_0=0.92, beta=0.96193, Eampfile=nan)
sscav = cavity(f=704.21e6, delF=nan, I_b0=50e-3, I_b=50e-3, V_cav=18.7927e6*0.97885161, beamph=-14.0,
        RoverQ=477, Ql=8.123e5, beta_0=0.92, beta=0.96193, Eampfile=nan)
#sscavperf = cavity(f=352.21e6, delF=nan, I_b0=50e-3, I_b=50e-3, V_cav=5.16749*0.58731097, beamph=-29.650191,
#        RoverQ=500, Ql=1.421e5, beta_0=0.46, beta=0.316718, Eampfile='Ezvszspoke.dat')
#sscav = cavity(f=352.21e6, delF=nan, I_b0=50e-3, I_b=50e-3, V_cav=5.16749e6*0.58731097, beamph=-29.650191,
#        RoverQ=500.0, Ql=1.421e5, beta_0=0.46, beta=0.316718, Eampfile='Ezvszspoke.dat')
print "Detuning = %f" % sscav.delF
print "QL = %e" % sscav.Ql
RoverQscaling = scale_roverq(beta_0=sscav.beta_0, file=sscav.Eampfile)(beta=sscav.beta, freq=sscav.f)[0]
print "RoverQscaling = %f" % RoverQscaling
cav = transientcavity(freq=sscav.f, QL=sscav.Ql, RoverQ=sscav.RoverQ, detuning=sscav.delF)
cavanal = analysis(tcavobj=cav, timestop=5e-3, acc=1e-3)

pulselength = 2.8e-3
faulttime = 1e-3
Ig, Ibeam, tinj = calcCurrents(cavobj=sscav, tcavobj=cav, 
        time=cavanal.time, pulselength=pulselength, RoverQscaling=RoverQscaling)
for i in range(len(cavanal.time)):
    if cavanal.time[i] < tinj + faulttime:
        continue
    Ibeam[i] = 0 + 0j
I = Ig + Ibeam
print "max(Ibeam) = %f" % max(abs(Ibeam))

print "Running..."
#profile.run('cavanal.run(iniV=0, I=I)')
cavanal.run(iniV=0, I=I)
print "Done!"

Vinc  = 0.5 * Ig * exp(cav.phi*1j) * cav.RL
Pinc  = abs(Vinc)**2 / cav.RL
Vref  = abs(Vinc) - sqrt(cavanal.V[:,0]**2 + cavanal.V[:,1]**2)
Pref  = Vref**2 / cav.RL
Vbeam = 0.5 * Ibeam * exp(cav.phi*1j) * cav.RL
Pbeam = abs(Vbeam)**2 / cav.RL

fig2 = plt.figure(2)
ax1fig2 = fig2.add_subplot(211)
ax1fig2.plot(1e3*cavanal.time, 1e-6*Vinc.real, 'b', label='Real incident voltage')
ax1fig2.plot(1e3*cavanal.time, 1e-6*Vinc.imag, 'r', label='Imag incident voltage')
ax1fig2.plot(1e3*cavanal.time, 1e-6*abs(Vbeam), 'k', label='Abs beam excited voltage')
ax1fig2.plot(1e3*cavanal.time, 1e-6*abs(Vinc), 'g', label='Abs incident voltage')
ax1fig2.set_ylabel('Voltage / MV')

ax2fig2 = fig2.add_subplot(212)
ax2fig2.plot(1e3*cavanal.time, abs(Vref)/abs(Vinc), label='Reflection coeff.')
ax2fig2.set_ylabel('Reflection coefficient')

fig = plt.figure(1)
ax1 = fig.add_subplot(211)
ax2 = fig.add_subplot(212)
ax1.plot(1e3*cavanal.time, 1e-6*cavanal.V[:,0], '-', label='Real cavity voltage')
ax1.plot(1e3*cavanal.time, 1e-6*cavanal.V[:,1], '-', label='Imag cavity voltage')
ax2.plot(1e3*cavanal.time, 1e-6*sqrt(cavanal.V[:,0]**2 + cavanal.V[:,1]**2), '-', label='Absolute cavity voltage')
ax2.plot(1e3*cavanal.time, 1e-6*Vref, 'r', label='Reflected voltage')

fig3 = plt.figure(3)
ax1fig3 = fig3.add_subplot(111)
ax1fig3.plot(1e3*cavanal.time, rad2deg(arctan(cavanal.V[:,1]/cavanal.V[:,0])), label='RF phase')
ax1fig3.set_ylabel('RF phase (degrees)')

fig4 = plt.figure(4)
ax1fig4 = fig4.add_subplot(111)
ax1fig4.plot(1e3*cavanal.time, Pref/1e3, label='Reflected: %0.3f J' % (sum(Pref * cavanal.timestep)))
ax1fig4.set_ylabel('Power (kW)')

for a in [ax1,ax2,ax1fig2,ax2fig2,ax1fig3,ax1fig4]:
    a.grid()
    a.set_xlabel('Time / ms')
    a.set_xlim(right=1e3*cavanal.time[-1])
    a.axvline(x=tinj*1e3, color='g')
    a.axvline(x=(tinj+pulselength)*1e3, color='g')
    a.legend(loc=0)
ax1fig2.legend(loc=1)
ax1.set_ylabel('Voltage / MV')
ax2.set_ylabel('Abs voltage / MV')

plt.show()

